import numpy as np
from scipy.interpolate import griddata  # 引入scipy中的二维插值库
import os
import matplotlib.pyplot as plt


def station2d_interpolate_to_grid(file):
    """
    func : 将站点数据插值到等经纬度格点
    inputs:
        lon: 站点的经度
        lat: 站点的纬度
        data: 对应经纬度站点的 气象要素值
        loc_range: [lat_min,lat_max,lon_min,lon_max]。站点数据插值到loc_range这个范围
        det_grid: 插值形成的网格空间分辨率
        method: 所选插值方法，默认 0.125
    return:

        [lon_grid,lat_grid,data_grid]
    """
    with open(file, encoding="GBK") as sf:
        sdls = sf.readlines()[4:]

    lonl = []
    latl = []
    indexes = []
    values = []
    rainmax = 0
    for sd in sdls:
        # try:
        sid, lon, lat, _, rain = sd.split()
        if latmin <= float(lat) <= latmax and lonmin <= float(lon) <= lonmax:
            x = round((latmax - float(lat)) / reso)
            y = round((float(lon) - lonmin) / reso)
            lonl.append(y)
            latl.append(x)
            values.append(float(rain))
            indexes.append([x, y])
            rainmax = max(float(rain), rainmax)
        # if float(rain)>0:
        #     print(rain)
        # else:
        #     # print("bad data", stfile, sd)
        #     pass
    # except:
    #     print("error found!")
    idd = (np.array(latl), np.array(lonl))

    data = np.array(values).reshape(-1, 1)

    # shape = [n,2]
    points = np.vstack(indexes)

    lon_grid, lat_grid = np.meshgrid(np.arange(ydim), np.arange(xdim))

    # step3:进行网格插值
    gdata = griddata(points, data, (lon_grid, lat_grid), method='cubic', fill_value=0).squeeze()
    gdata[gdata < 0] = 0
    rainmax *= 1.2
    gdata[gdata > rainmax] = rainmax

    # plt.subplot(222)
    # plt.imshow(gdata, extent=(0, 10, 0, 10), origin='lower')
    plt.plot(points[:, 0], points[:, 1], 'k.', ms=1)
    plt.title('cubic')
    plt.gcf().set_size_inches(8, 8)
    plt.show()
    radio = np.zeros_like(gdata)
    for idx in indexes:
        x, y = idx
        for i in range(5):
            for j in range(5):
                st_idxs = []
                if 0 <= x - i < radio.shape[0] and 0 <= y - j < radio.shape[1]:
                    st_idxs.append((x - i, y - j))
                if 0 <= x - i < radio.shape[0] and 0 <= y + j < radio.shape[1]:
                    st_idxs.append((x - i, y + j))
                if 0 <= x + i < radio.shape[0] and 0 <= y - j < radio.shape[1]:
                    st_idxs.append((x + i, y - j))
                if 0 <= x + i < radio.shape[0] and 0 <= y + j < radio.shape[1]:
                    st_idxs.append((x + i, y + j))
                for sidx in list(set(st_idxs)):
                    radio[sidx[0]][sidx[1]] = 1 / (((i + 1) ** 2 + (j + 1) ** 2) / 2) ** 0.5

    return gdata, radio


if __name__ == '__main__':
    # 实际选取区域: 19.01, 27.8, 102.2, 113.79  [120:-200, 200:-160]
    latmin, latmax, lonmin, lonmax = 21.01, 26.6, 104.2, 112.19
    xdim, ydim, reso = 560, 800, 0.01
    # stfile = "data/202106142024.000"
    sdir = "gridData"
    import glob
    import os

    files = glob.glob("data/20211010/*.000")[:5]
    for stfile in files:
        gdata, radio = station2d_interpolate_to_grid(stfile)

        gfile = os.path.split(stfile)[-1].split('.')[0]
        # np.save(os.path.join(sdir, gfile), gdata)
        pass
